#!coding:utf-8
import numpy as np
import json
import rospy
from pupper.config import Command, State, BehaviorState, Configuration
from pupper.controller import Controller
from pupper.inverse_kinematics import four_legs_inverse_kinematics
from pupper.hardware_interface import HardwareInterface
from pi_driver.srv import SetString, SetStringResponse, GetString, GetStringResponse

BehaviorStateMap = {
    0: BehaviorState.REST,
    1: BehaviorState.TROT,
    5: BehaviorState.DEACTIVATED,
}


def limit_to(value, max_value):
    if value > max_value:
        value = max_value
    elif value < -max_value:
        value = -max_value
    return value


class PupperDriverNode:
    def __init__(self):
        self.node_name = rospy.get_name()
        rospy.loginfo("[%s] Initializing......" % (self.node_name))
        # 初始化各种对象
        self.config = Configuration()
        self.controller = Controller(
            self.config,
        )
        self.state = State()
        self.command = Command()
        self.hardware_interface = HardwareInterface()

        # 初始姿态
        self.state.foot_locations = (
            self.config.default_stance
            + np.array([0, 0, self.command.height])[:, np.newaxis]
        )
        # 使用 ros 定时器, 避免手动计算
        rospy.Timer(rospy.Duration(1.0/100), self.step)
        rospy.Service('/ubiquityrobot/pupper_driver_node/set_command',
                      SetString, self.cbSetCommand)
        rospy.Service('/ubiquityrobot/pupper_driver_node/get_state',
                      GetString, self.cbGetState)
        rospy.Service('/ubiquityrobot/pupper_driver_node/set_servo_angles',
                      SetString, self.cbSetServoAngles)
        rospy.Service('/ubiquityrobot/pupper_driver_node/set_foot_locations',
                      SetString, self.cbSetFootLocations)
        rospy.loginfo("[%s] Initialized." % (self.node_name))

    def step(self, event=None):
        if self.state.behavior_state == BehaviorState.DEACTIVATED:
            return
        # 执行一次坐标计算
        self.controller.run(self.state, self.command)

        # 坐标到关节角度
        joint_angles = four_legs_inverse_kinematics(
            self.state.rotated_foot_locations, self.config)

        # 关节角度到舵机指令
        self.hardware_interface.set_actuator_postions(joint_angles)

    def cbSetCommand(self, params):
        try:
            data = json.loads(params.data)
            print(data)
            if 'horizontal_velocity_x' in data:
                value = limit_to(data['horizontal_velocity_x'],
                                 self.config.max_x_velocity)
                self.command.horizontal_velocity[0] = value
            if 'horizontal_velocity_y' in data:
                value = limit_to(data['horizontal_velocity_y'],
                                 self.config.max_y_velocity)
                self.command.horizontal_velocity[1] = value
            if 'roll' in data:
                self.command.roll = limit_to(
                    data['roll'], self.config.max_roll)
            if 'pitch' in data:
                self.command.pitch = limit_to(
                    data['pitch'], self.config.max_pitch)
            if 'yaw_rate' in data:
                self.command.yaw_rate = limit_to(
                    data['yaw_rate'], self.config.max_yaw_rate)
            if 'height' in data:
                if data['height'] < 0.08:
                    data['height'] = 0.08
                if data['height'] > 0.22:
                    data['height'] = 0.22
                self.command.height = -data['height']
            if 'behavior_state' in data and data['behavior_state'] in BehaviorStateMap:
                self.state.behavior_state = BehaviorStateMap[data['behavior_state']]
            if 'servo_speed' in data and data['servo_speed'] > 0 and data['servo_speed'] <= 2000:
                self.hardware_interface.speed = int(data['servo_speed'])
        except Exception as e:
            print(e)
        return SetStringResponse(params.data)

    def get_command(self):
        return self.command

    def cbGetState(self, params):
        state = {
            "servo_angles": self.hardware_interface.servo_angles,
            "foot_locations": self.state.foot_locations.tolist()
        }
        return GetStringResponse(json.dumps(state))

    def cbSetServoAngles(self, params):
        joint_angles = json.loads(params.data)
        self.hardware_interface.set_actuator_angles(
            joint_angles)
        return SetStringResponse()

    def cbSetFootLocations(self, params):
        foot_locations = np.array(json.loads(params.data))
        print(foot_locations)
        joint_states = four_legs_inverse_kinematics(
            foot_locations, self.controller.config
        )
        self.hardware_interface.set_actuator_postions(
            joint_states)
        return SetStringResponse()

    def onShutdown(self):
        # 节点关闭时要执行的操作
        rospy.loginfo("[%s] Closing ." % (self.node_name))
        rospy.loginfo("[%s] Shutdown." % (self.node_name))


if __name__ == '__main__':
    rospy.init_node('pupper_driver_node')
    node = PupperDriverNode()
    rospy.on_shutdown(node.onShutdown)
    rospy.spin()
